package com.tianji.aigc.memory;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.entity.RedisMessage;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 基于Redis的会话记忆处理类
 *
 * @Author mr.wu
 * @Date 2025-9-2 17:21
 */
@Component
public class RedisChatMemory implements ChatMemory {

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 对话完成添加消息
     * @param conversationId  对话ID
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        //1.拼接RedisKey
        String key = getRedisKey(conversationId);

        //2.将消息转为字符串类型
        List<String> messageStrList = messages.stream().map(x -> {
            RedisMessage redisMessage = BeanUtil.toBean(x, RedisMessage.class);
            redisMessage.setTextContent(x.getText());
            if(x instanceof AssistantMessage assistantMessage) {
                redisMessage.setToolCalls(assistantMessage.getToolCalls());
            } else if(x instanceof ToolResponseMessage toolResponseMessage) {
                redisMessage.setToolResponses(toolResponseMessage.getResponses());
            }
            return JSONUtil.toJsonStr(redisMessage);
        }).toList();

        //3.将消息存入Redis的List中
        redisTemplate.opsForList().leftPushAll(key,messageStrList);
    }

    /**
     * 获取对话ID相关的消息列表
     * @param conversationId  对话ID
     * @param lastN  消息数量
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        //1.拼接RedisKey
        String key = getRedisKey(conversationId);

        //2.获取消息字符串列表
        List<String> messageStrList = redisTemplate.opsForList().range(key, 0, lastN);

        //3.将消息字符串转为消息对象
        List<Message> messageList = messageStrList.stream().map(x -> {
            RedisMessage redisMessage = JSONUtil.toBean(x, RedisMessage.class);

            Message message = null;
            if(redisMessage.getMessageType().equals(MessageType.ASSISTANT)) {
                message = new AssistantMessage(redisMessage.getTextContent(),redisMessage.getMetadata(),redisMessage.getToolCalls(),redisMessage.getMedia());
            } else if(redisMessage.getMessageType().equals(MessageType.TOOL)) {
                message = new ToolResponseMessage(redisMessage.getToolResponses(),redisMessage.getMetadata());
            } else if(redisMessage.getMessageType().equals(MessageType.USER)) {
                message = new UserMessage(redisMessage.getMessageType(), redisMessage.getTextContent(),redisMessage.getMedia(),redisMessage.getMetadata());
            } else {
                message = new SystemMessage(redisMessage.getTextContent());
            }

            return message;
        }).toList();
        return messageList;
    }

    /**
     * 清理对话消息
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        //1.拼接RedisKey
        String key = getRedisKey(conversationId);

        //2.删除缓存中所有消息
        redisTemplate.unlink(key);
    }

    /**
     * 拼接RedisKey
     * @param conversationId
     * @return
     */
    private String getRedisKey(String conversationId){
        return "CHAT:" + conversationId;
    }
}
